import streamlit as st
import pandas as pd
from io import BytesIO
import sys
from contextlib import redirect_stdout, redirect_stderr
import io
import traceback
from code_interpreter1 import CodeInterpreter

# 初始化代码解释器
if 'interpreter' not in st.session_state:
    st.session_state.interpreter = CodeInterpreter()

# 设置页面标题
st.title("Excel 处理工具 & 代码解释器")

# 上传文件
uploaded_file = st.file_uploader("上传 Excel 文件", type=["xlsx", "xls"])

# 全局变量用于存储DataFrame
if 'df' not in st.session_state:
    st.session_state.df = None

if uploaded_file is not None:
    # 读取 Excel 文件
    st.session_state.df = pd.read_excel(uploaded_file)

    # 显示原始数据
    st.subheader("原始数据")
    st.write(st.session_state.df)

    # 代码编辑器
    st.subheader("Python 代码编辑器")
    default_code = """# 在这里编写处理Excel的Python代码
# 使用变量 df 来引用上传的Excel数据
# 支持的模块: pandas, numpy, math, random, json, re, time, copy
# 示例:
# df['新列'] = df['某列'].apply(lambda x: x.upper())
"""
    code = st.text_area("编辑代码", value=default_code, height=300, key="code_editor")

    # 创建一个包含运行按钮的列布局
    col1, col2 = st.columns([6, 1])
    with col2:
        run_button = st.button("运行代码", type="primary")

    if run_button and code:
        st.subheader("执行结果")

        # 使用代码解释器执行代码
        result = st.session_state.interpreter.execute_code(code, st.session_state.df)

        if result['success']:
            # 更新DataFrame
            st.session_state.df = result['df']

            # 显示代码输出
            if result['output']:
                st.text("代码输出:")
                st.code(result['output'])

            # 显示处理后的数据
            st.subheader("处理后的数据")
            st.write(st.session_state.df)

            # 提供下载处理后的文件选项
            output_excel = BytesIO()
            with pd.ExcelWriter(output_excel, engine='openpyxl') as writer:
                st.session_state.df.to_excel(writer, index=False)
            output_excel.seek(0)

            st.download_button(
                label="下载处理后的 Excel 文件",
                data=output_excel,
                file_name="processed_excel.xlsx",
                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
            )
        else:
            st.error("执行代码时出错:")
            st.code(result['error'])